{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run DESYNC-MHP MLE on a toy example\n",
    "\n",
    "This notebook shows an example of the DESYNC-MLE algorithm on a small toy network.\n",
    "The toy example is the one used in the paper and is defined as follows: \n",
    "\n",
    "We generate a synthetic realization from a multivariate Hawkes process with exponential excitation functions in $d=2$ dimensions (i.e., processes $N_1$ and $N_2$). Process $N_1$ excites itself and influences process $N_2$ (but not the other way around). In other words, the excitation matrix is \n",
    "\n",
    "\\begin{align}\n",
    "    A = \n",
    "    \\begin{bmatrix} \n",
    "    \\alpha_{11}  e^{-\\beta t} 1_{\\{t > 0\\}} & 0 \\\\\n",
    "    \\alpha_{21}  e^{-\\beta t} 1_{\\{t > 0\\}} & 0 \\\\\n",
    "    \\end{bmatrix}  \n",
    "\\end{align}\n",
    "\n",
    "Therefore, the intensity of the processes are\n",
    "\\begin{align}\n",
    " \\lambda_1(t) = \\mu_1 + \\alpha_{11} e^{-\\beta t} 1_{\\{t > 0\\}}  + \\alpha_{21} e^{-\\beta t} 1_{\\{t > 0\\}} \n",
    " \\quad \\text{and} \\quad\n",
    " \\lambda_2(t) = \\mu_2.\n",
    "\\end{align}\n",
    "\n",
    "As discussed in the paper, for a synchronization noise such that $z_1 > z_2$, the cause and effect arrivals start to blur and the classic maxmimum likelihood estimation converges to a complete graph. However, our approach, DESYNC-MHP MLE, is able to accurately recover the true causal structure of the excitation matrix."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "Load libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "if '..' not in sys.path:\n",
    "    # Allows to import the `lib` module holding the code for DESYNC-MHP MLE\n",
    "    sys.path.append('..')\n",
    "\n",
    "# External libs\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "\n",
    "# Use `tick` to simulate synthetic data and for the baseline `Classic MLE` estimator\n",
    "from tick.hawkes.simulation import SimuHawkesExpKernels\n",
    "from tick.hawkes.inference import HawkesADM4\n",
    "\n",
    "# Code for the `DESYNC-MLE` algorithm\n",
    "from lib.inference.learner_conditional_mle import HawkesExpKernConditionalMLE\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "# Graph drawing parameters\n",
    "nx_draw_params = {\n",
    "    'pos':{0: (0,0), 1: (1, 0)}, \n",
    "    'node_size': 2e3, 'node_color': 'w', \n",
    "    'labels': {0: '$N_1$', 1: '$N_2$'},\n",
    "    'edgecolors':'k', 'arrowsize': 50, 'width': 2,\n",
    "    'font_size': 20\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize random seed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(12345678)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## 1. Define parameters\n",
    "\n",
    "Set parameters for the simulation of synthetic data. \n",
    "\n",
    "* `n_nodes` is the number of nodes of the process,\n",
    "* `decay` is the exponential decay $\\beta$ of the kernels,\n",
    "* `end_time` is the length of the observation window,\n",
    "* `noise_scale` is the variance used to sample the synchronization noise\n",
    "* `baseline` is the vector of background intensity $\\boldsymbol{\\mu} = [\\mu_1, \\mu_2]$,\n",
    "* `adjacency` is the matrix of magnitudes $\\{\\alpha_{ij}\\}$ of the kernels.\n",
    "\n",
    "Feel free to try the code with other parameters (*e.g.* larger network, increasing the noise scale, ...)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "n_nodes = 2         # Number of nodes\n",
    "decay = 1.0         # Exponential decay\n",
    "end_time = 100e3    # Length of the observation window\n",
    "noise_scale = 1.0   # Scale (variance) of the noise \n",
    "\n",
    "# Background intensity\n",
    "mu = 0.05\n",
    "baseline = mu * np.ones(n_nodes, dtype='float')\n",
    "\n",
    "# Excitation matrix\n",
    "adjacency = np.array([[0.9, 0.0],\n",
    "                      [0.9, 0.0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print('True baseline:')\n",
    "print(baseline.round(2))\n",
    "print('True adjacency:')\n",
    "print(adjacency.round(2))\n",
    "\n",
    "# Draw the graph\n",
    "print('\\nEstimated excitation graph:')\n",
    "graph_true = nx.DiGraph(adjacency.T)\n",
    "plt.figure(figsize=(8, 1))\n",
    "nx.draw_networkx(graph_true, **nx_draw_params)\n",
    "plt.axis('off');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## 2. Simulate data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sample a random synchronization noise value for each dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z_true = np.random.normal(scale=noise_scale, size=n_nodes)\n",
    "\n",
    "# Order the noise value to be in the \"bad\" regime \n",
    "# (where the events from N_1 are delayed more than those from N_2)\n",
    "z_true = np.sort(z_true)[::-1]\n",
    "\n",
    "z_true"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sample a realization of the multivariate Hawkes process with exponential excitation kernels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Compute the observation window adjusted to the noise to avoid edge effects\n",
    "adjusted_end_time = end_time + z_true.min() - z_true.max()\n",
    "adjusted_start_time = z_true.max()\n",
    "\n",
    "seed = 42\n",
    "\n",
    "simu = SimuHawkesExpKernels(adjacency=adjacency, decays=decay, baseline=baseline,\n",
    "                            end_time=adjusted_end_time, seed=42)\n",
    "simu.simulate()\n",
    "\n",
    "print('\\nNumber of events simulated:', list(map(len, simu.timestamps)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add synchronization noise to the observed events."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hold the noisy realizations\n",
    "noisy_events = []\n",
    "\n",
    "# For each dimensions\n",
    "for m, orig_events_m in enumerate(simu.timestamps):\n",
    "    # Add the noise\n",
    "    events_m = orig_events_m + z_true[m] - adjusted_start_time\n",
    "    # Filter the events outside the observation window to avoid edge effects\n",
    "    events_m = events_m[(events_m >= 0) & (events_m < end_time)]\n",
    "    noisy_events.append(events_m)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the observations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Original noiseless events:')\n",
    "for m, orig_events_m in enumerate(simu.timestamps):\n",
    "    print(f'{m:>d} | {orig_events_m}')\n",
    "\n",
    "print(f'\\nNoise value: {z_true}\\n')\n",
    "    \n",
    "print('Noisy events:')\n",
    "for m, events_m in enumerate(noisy_events):\n",
    "    print(f'{m:>d} | {events_m}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## 3. Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the starting parameters values at random."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z_start = np.random.normal(scale=noise_scale, size=n_nodes)\n",
    "\n",
    "baseline_start = np.random.uniform(0.0, 1.0, size=n_nodes)\n",
    "adjacency_start = np.random.uniform(0.0, 1.0, size=(n_nodes, n_nodes)) \n",
    "\n",
    "# Stack all parameters into a single vector\n",
    "theta_start = np.hstack((baseline_start, adjacency_start.ravel()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "###  3.1. Classic MLE (baseline method)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "naive_learner = HawkesADM4(decay, lasso_nuclear_ratio=1.0,\n",
    "                           verbose=True, print_every=10, record_every=10)\n",
    "naive_learner.fit(noisy_events, end_time,\n",
    "                  baseline_start=baseline_start,\n",
    "                  adjacency_start=adjacency_start);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('\\nEstimated baseline:')\n",
    "print(naive_learner.baseline.round(2))\n",
    "print('Ground truth baseline:')\n",
    "print(baseline.round(2))\n",
    "\n",
    "print('\\nEstimated excitation matrix:')\n",
    "print(naive_learner.adjacency.round(2))\n",
    "print('Ground truth excitation matrix:')\n",
    "print(adjacency.round(2))\n",
    "\n",
    "adj_err = (naive_learner.adjacency > 0.05).astype(int) - (adjacency > 0).astype(int)\n",
    "print('\\nErrors in the excitation matrix: ({:d} errors total) (FP: 1, FN: -1)'.format(np.abs(adj_err).sum()))\n",
    "print(adj_err)\n",
    "\n",
    "print('\\nEstimated excitation graph:')\n",
    "plt.figure(figsize=(8, 1))\n",
    "naive_learner.adjacency[naive_learner.adjacency < 0.05] = 0\n",
    "graph_classic_mle = nx.DiGraph(naive_learner.adjacency.T)\n",
    "nx.draw_networkx(graph_classic_mle, **nx_draw_params)\n",
    "plt.axis('off');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "### 3.2. Run DESYNC-MHP MLE (contribution)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initialize the learner object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the step size\n",
    "step_size_func = lambda t: 3.0 * (t+1) ** -0.5\n",
    "\n",
    "learner = HawkesExpKernConditionalMLE(decay,\n",
    "    approx_type='smooth',      # Use the smooth approximation of the exponential kernel\n",
    "    decay_neg=50.0,            # Negative decay of the smooth kernel\n",
    "    gamma=100.0,               # Transition rate between negative/positive parts of the smooth kernel\n",
    "    n_threads=2,               # Number of threads used for computations\n",
    "    max_iter=1000,             # Maximum number of iterations\n",
    "    step_z=step_size_func,     # Step size\n",
    "    step_theta=step_size_func, # Step size\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run the optimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "coeffs = learner.fit(\n",
    "    noisy_events, end_time,   # Set observed data\n",
    "    z_start=z_start,          # Initial noise parameters value\n",
    "    theta_start=theta_start,  # Initial MHP parameters\n",
    "    seed=42,                  # Set random seed for SGD\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Visualize the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('\\nEstimated baseline:')\n",
    "print(learner.baseline.round(2))\n",
    "print('Ground truth baseline:')\n",
    "print(baseline.round(2))\n",
    "\n",
    "print('\\nEstimated excitation matrix:')\n",
    "print(learner.adjacency.round(2))\n",
    "print('Ground truth excitation matrix:')\n",
    "print(adjacency.round(2))\n",
    "\n",
    "adj_err = (learner.adjacency > 0.05).astype(int) - (adjacency > 0).astype(int)\n",
    "print('\\nErrors in the excitation matrix: ({:d} errors total) (FP: 1, FN: -1)'.format(np.abs(adj_err).sum()))\n",
    "print(adj_err)\n",
    "\n",
    "print('\\nEstimated excitation graph:')\n",
    "plt.figure(figsize=(8, 1))\n",
    "learner.adjacency[learner.adjacency < 0.05] = 0\n",
    "graph = nx.DiGraph(learner.adjacency.T)\n",
    "nx.draw_networkx(graph, **nx_draw_params)\n",
    "plt.axis('off');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
